package com.ddbq.pytorchmode.service;

import ai.djl.MalformedModelException;
import ai.djl.Model;
import ai.djl.inference.Predictor;
import ai.djl.modality.Classifications;
import ai.djl.modality.cv.Image;
import ai.djl.modality.cv.ImageFactory;
import ai.djl.modality.cv.transform.CenterCrop;
import ai.djl.modality.cv.transform.Normalize;
import ai.djl.modality.cv.transform.Resize;
import ai.djl.modality.cv.transform.ToTensor;
import ai.djl.modality.cv.translator.ImageClassificationTranslator;
import ai.djl.repository.zoo.Criteria;
import ai.djl.repository.zoo.ModelNotFoundException;
import ai.djl.repository.zoo.ModelZoo;
import ai.djl.training.util.ProgressBar;
import ai.djl.translate.Pipeline;
import ai.djl.translate.TranslateException;
import ai.djl.translate.Translator;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
/**
 * @program: pytorch-mode
 * @description:
 * @author: zqs
 * @date： 2024-07-23-10:29
 */
@Service
public class PytorchService {

    public void predict(){
        Pipeline pipeline = new Pipeline();
        pipeline.add(new Resize(256))
                .add(new CenterCrop(224, 224))
                .add(new ToTensor())
                .add(new Normalize(
                        new float[]{0.485f, 0.456f, 0.406f},
                        new float[]{0.229f, 0.224f, 0.225f}));

        Translator<Image, Classifications> translator = ImageClassificationTranslator.builder()
                .setPipeline(pipeline)
                .optApplySoftmax(true)
                .build();

        System.setProperty("ai.djl.repository.zoo.location", "src/main/resources/yolo");
        Criteria<Image, Classifications> criteria = Criteria.builder()
                .setTypes(Image.class, Classifications.class)
                // only search the model in local directory
                // "ai.djl.localmodelzoo:{name of the model}"
                .optArtifactId("ai.djl.localmodelzoo:yolo")
                .optTranslator(translator)
                .optProgress(new ProgressBar()).build();
        try {
            Model model = ModelZoo.loadModel(criteria);
            // 自己本地
            File fs=new File("C:\\Users\\zouqi\\Pictures\\AI\\e8e4-f7a99e8e44ef84bbf63c5e98513149a8.jpg");
            Image img = ImageFactory.getInstance().fromInputStream(new FileInputStream(fs));
            Predictor<Image, Classifications> predictor = model.newPredictor(translator);
            Classifications classifications = predictor.predict(img);
            System.out.println(classifications);
        } catch (IOException | ModelNotFoundException | MalformedModelException | TranslateException e) {
            e.printStackTrace();
        }


    }
}
